# -*- coding: utf-8 -*-
##########
# scan code based on Red9 Consultancy script
##########

'''
    Malware fixer for a known Maya virus discovered in a client file, and were
    subsequently infected with for a few days. The Malware is basically a scriptNode 
    in a scene which auto-generates a userSetup.mel, various mel globals and a scriptJob.
    These are all designed to propagate the scriptNodes in all future files worked on
    within an infected machine.
    
    The ultimate aim of this malware seems to be to hang all Maya files and systems on, or after the 26/06/2020

    Original Author: Mark Jackson
    email: info@red9consultancy.com
    Red9 : http://red9consultancy.com
    Red9 Vimeo : https://vimeo.com/user9491246
    
    to run import this module and run
    clean_malware()

'''

import sys
import os
import logging
import logging.handlers
import tempfile

import maya.cmds as cmds
import maya.mel as mel

from stat import S_IWUSR, S_IREAD


# create a log file of found issues
def MayaScannerLogFile():
    return os.path.join(tempfile.gettempdir(), "MayaScannerLog.txt")


logging.basicConfig()
log = logging.getLogger('Autodesk.MayaScanner')

if not len(log.handlers):
    filehandler = logging.handlers.RotatingFileHandler(MayaScannerLogFile(), mode='a', backupCount=1, delay=True)
    log.addHandler(filehandler)


# roll over the log file
def rollOverLogFile():
    log.handlers[0].doRollover()


# from Log import log
# ask user if ok to attempt to fix the issue

def userConfirmFix(msgString, smode=0):
    if 'state' not in userConfirmFix.__dict__:
        userConfirmFix.state = 0
    if smode == 1:
        userConfirmFix.state = 0
        return 0

    if userConfirmFix.state == 0:
        # check with user to continue with attempted fix
        answer = cmds.confirmDialog(title=msgString, message='Attempt to fix issue?',
                                    button=['Yes', 'No'], defaultButton='Yes', cancelButton='No', dismissString='No')
        if cmds.about(batch=True) or answer == 'Yes':
            userConfirmFix.state = 1

    # return status of user choice 
    return userConfirmFix.state


def reportIssue(msgString, smode=0):
    if 'cnt' not in reportIssue.__dict__:
        reportIssue.cnt = 0
    if smode == 1:
        reportIssue.cnt = 0
        return
    first = (reportIssue.cnt == 0)
    reportIssue.cnt = 1

    if first:
        log.info(u"正在检查文件中的问题: %s" % cmds.file(q=True, sn=True))

    log.info(msgString)


def test_userSetup():
    '''
    test if the userSetup.mel has been created, or appended by the malware
    '''
    prefs = os.path.dirname(os.path.dirname(cmds.about(env=True)))
    usersetup = os.path.normpath(os.path.join(prefs, 'scripts', 'userSetup.mel'))
    status = False
    if os.path.exists(usersetup):
        f = open(usersetup, "r")
        data = f.read()

        if 'fuck_All_U' in data:
            status = 'compromised'
            reportIssue(u'userSetup.mel : 已被恶意破坏!')
        if all([len(data) >= 4118,
                '// Maya Mel UI Configuration File.Maya Mel UI Configuration File..\n// \n//\n//  This script is machine generated.  Edit at your own risk' in data,
                'string $chengxu' in data]):
            reportIssue(u'userSetup.mel : 这是恶意创建的文件!')
            status = 'delete'
    return usersetup, status


def fix_userSetup(prefixTitle):
    issueFound = 0
    issueFixed = 0
    usersetup, status = test_userSetup()
    if status == 'delete':
        # fouund an issue notifiy user, ask to attempt fix
        if userConfirmFix('Autodesk.MayaScanner: %s : found corrupted userSetup.mel' % prefixTitle):
            #
            os.chmod(usersetup, S_IWUSR | S_IREAD)
            os.remove(usersetup)
            reportIssue('Removed : userSetup.mel: %s' % usersetup)
            issueFixed = 1
        # we still want to know of this for fail report
        issueFound = 1
    return issueFound, issueFixed


def test_scriptNodes():
    '''
    test for the known scriptNode with specific data
    '''
    malware_scripts = []
    for script in cmds.ls(type='script'):
        if 'MayaMelUIConfigurationFile' in script.split('|')[-1].split(':')[-1]:
            scriptdata = cmds.scriptNode(script, bs=True, q=True)
            if 'This script is machine generated.  Edit at your own risk' in scriptdata and 'fuck_All_U' in scriptdata:
                malware_scripts.append(script)
                reportIssue(u'发现scriptNode 有存在恶意节点 : %s' % script)
    return malware_scripts


def fix_scriptNodes(prefixTitle):
    issueFound = 0
    issueFixed = 0
    for script in test_scriptNodes():
        if userConfirmFix('Autodesk.MayaScanner: %s : found corrupted scriptNode' % prefixTitle):
            cmds.delete(script)
            reportIssue('Removed : scriptNode: %s' % script)
            issueFixed += 1
        issueFound += 1
    return issueFound, issueFixed


def test_scriptJob():
    '''
    test for the scriptJob this Malware geenerates
    '''
    scriptjob_id = None
    try:
        a = -1
        if (mel.eval('whatIs "$autoUpdateAttrEd_aoto_int"') != "Unknown"):
            a = int(mel.eval('$temp=$autoUpdateAttrEd_aoto_int'))
        for job in cmds.scriptJob(listJobs=True):
            if job.startswith(str(a)):
                scriptjob_id = a
                reportIssue(u'发现被损坏的恶意节点 : scriptJob : %s' % a)
                break
    except:
        reportIssue(u'scriptjob 没有发现被损坏的恶意节点')
    return scriptjob_id


def fix_scriptJob(prefixTitle):
    issueFound = 0
    issueFixed = 0
    id = test_scriptJob()
    if id:
        if userConfirmFix('Autodesk.MayaScanner: %s : found corrupted scriptJob' % prefixTitle):
            cmds.scriptJob(kill=id, force=True)
            reportIssue('Removed : scriptJob ID: %s' % id)
            issueFixed = 1
        # return 1 indicating found script job to kill
        issueFound = 1
    return issueFound, issueFixed


#
# script to run, would like to know if running single file or multiple for file logging
def clean_malware(prefixTitle):
    '''
    remove the damn thing and all nodes created!
    '''

    # first kill the mel globals!
    mel_globals = ['UI_Mel_Configuration_think',
                   'UI_Mel_Configuration_think_a',
                   'UI_Mel_Configuration_think_b',
                   'autoUpdateAttrEd_SelectSystem',
                   'autoUpdatcAttrEd',
                   'autoUpdatoAttrEnd']
    for glb in mel_globals:
        if mel.eval('whatIs("%s")' % glb) == 'Mel procedure entered interactively.':
            mel.eval('global proc %s(){error -sl "attempted to run corrupted command: %s";}' % (glb, glb))

    # run the base fixes
    issuesFound = 0
    issuesFixed = 0

    sJobFound, sJobFixed = fix_scriptJob(prefixTitle)
    sNodeFound, sNodeFixed = fix_scriptNodes(prefixTitle)
    sSetupFound, sSetupFixed = fix_userSetup(prefixTitle)

    issuesFound = sJobFound + sNodeFound + sSetupFound
    issuesFixed = sJobFixed + sNodeFixed

    return issuesFound, issuesFixed


# -
# ==========================================================================
# Copyright (C) 2020 Autodesk, Inc. and/or its licensors.  All 
# rights reserved.
#
# The coded instructions, statements, computer programs, and/or related 
# material (collectively the "Data") in these files contain unpublished 
# information proprietary to Autodesk, Inc. ("Autodesk") and/or its 
# licensors, which is protected by U.S. and Canadian federal copyright 
# law and by international treaties.
#
# The Data is provided for use exclusively by You. You have the right 
# to use, modify, and incorporate this Data into other products for 
# purposes authorized by the Autodesk software license agreement, 
# without fee.
#
# The copyright notices in the Software and this entire statement, 
# including the above license grant, this restriction and the 
# following disclaimer, must be included in all copies of the 
# Software, in whole or in part, and all derivative works of 
# the Software, unless such copies or derivative works are solely 
# in the form of machine-executable object code generated by a 
# source language processor.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. 
# AUTODESK DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED 
# WARRANTIES INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF 
# NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR 
# PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE, OR 
# TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS LICENSORS 
# BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL, 
# DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK 
# AND/OR ITS LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY 
# OR PROBABILITY OF SUCH DAMAGES.
#
# ==========================================================================
# +


# Jarry silence fix

def jarry_fix_scriptJob(fix_mode=True):
    issueFound = 0
    issueFixed = 0
    id = test_scriptJob()

    if id:
        if fix_mode:
            cmds.scriptJob(kill=id, force=True)
            reportIssue(u'已删除恶意节点 : scriptJob ID: %s' % id)
            issueFixed = 1
        # return 1 indicating found script job to kill
        issueFound = 1
    return issueFound, issueFixed


def jarry_fix_scriptNodes(fix_mode=True):
    issueFound = 0
    issueFixed = 0
    for script in test_scriptNodes():
        if fix_mode:
            cmds.delete(script)
            reportIssue(u'已删除恶意节点 : scriptNode: %s' % script)
            issueFixed += 1
        issueFound += 1
    return issueFound, issueFixed


def jarry_fix_userSetup(fix_mode=True):
    issueFound = 0
    issueFixed = 0
    usersetup, status = test_userSetup()
    if status == 'delete':
        # fouund an issue notifiy user, ask to attempt fix
        if fix_mode:
            #
            os.chmod(usersetup, S_IWUSR | S_IREAD)
            os.remove(usersetup)
            reportIssue(u'已删除恶意创建的配置文件 : userSetup.mel: %s' % usersetup)
            issueFixed = 1
        # we still want to know of this for fail report
        issueFound = 1
    return issueFound, issueFixed


def jarry_clean_malware(fix_mode=True):
    '''
    remove the damn thing and all nodes created!
    '''

    # first kill the mel globals!
    mel_globals = ['UI_Mel_Configuration_think',
                   'UI_Mel_Configuration_think_a',
                   'UI_Mel_Configuration_think_b',
                   'autoUpdateAttrEd_SelectSystem',
                   'autoUpdatcAttrEd',
                   'autoUpdatoAttrEnd']
    for glb in mel_globals:
        if mel.eval('whatIs("%s")' % glb) == 'Mel procedure entered interactively.':
            mel.eval('global proc %s(){error -sl "attempted to run corrupted command: %s";}' % (glb, glb))

    # run the base fixes
    issuesFound = 0
    issuesFixed = 0

    sJobFound, sJobFixed = jarry_fix_scriptJob(fix_mode)
    sNodeFound, sNodeFixed = jarry_fix_scriptNodes(fix_mode)
    sSetupFound, sSetupFixed = jarry_fix_userSetup(fix_mode)

    issuesFound = sJobFound + sNodeFound + sSetupFound
    issuesFixed = sJobFixed + sNodeFixed

    return issuesFound, issuesFixed


def jarry_currentfix_save(save=True):
    issuesFound, issuesFixed = jarry_clean_malware()
    if issuesFixed > 0:
        log.info(u'发现损坏 (%s),已经修复损坏 (%s)' % (issuesFound, issuesFixed))
        cmds.file(save=True, force=True)
        if save:
            s = cmds.file(q=1, sceneName=1)
            log.info(u'%s 文件修复后已保存!' % s)
    if issuesFound == 0:
        log.info(u'没有发现恶意损坏，文件跳过!')


def jarry_fileOpen_fix(fileName):
    if not os.path.exists(fileName):
        log.warning(u"%s 这个文件不存在" % fileName)
        return
    cmds.file(fileName, open=True, force=True, esn=0)
    jarry_currentfix_save()


def jarry_fileList_fix(fileList):
    if not fileList:
        log.warning(u'需要输入批量文件列表')
        return

    for fn in fileList:
        jarry_fileOpen_fix(fn)


'''
import MayaScannerCleaner as msc
reload(msc)
msc.jarry_fileOpen_fix(r'D:/putian/batch/c_GodTurtle_mod_v001.0001_source.mb')
'''
